package proxy

import (
	"bytes"
	"fmt"
	"gateway-server/internal/types"
	"io"
	"net/http"
)

type ModelServiceProxyImpl struct {
	ConnPool         ModelServiceConnPool
	ModelPathPrefix  string
	UniformedApiPath string
}

func NewModelServiceProxyImpl(connPool ModelServiceConnPool, modelPathPrefix string, uniformedApiPath string) ModelServiceProxy {
	return &ModelServiceProxyImpl{
		ConnPool:         connPool,
		ModelPathPrefix:  modelPathPrefix,
		UniformedApiPath: uniformedApiPath,
	}
}

func (msp *ModelServiceProxyImpl) GenerateRequest(apiKey string, originalRequestBody []byte, originalRequest *http.Request) (*http.Request, error) {
	backendURL := msp.ModelPathPrefix + msp.UniformedApiPath
	newReq, err := http.NewRequest(originalRequest.Method, backendURL, bytes.NewBuffer(originalRequestBody))
	if err != nil {
		return nil, fmt.Errorf("failed to create new request: %v", err)
	}
	newReq.Header = originalRequest.Header.Clone()
	newReq.Header.Set(types.AuthorizationHeader, types.AuthorizationPrefix+apiKey)
	return newReq, nil
}
func (msp *ModelServiceProxyImpl) ForwardRequest(req *http.Request) (*http.Response, error) {
	client, err := msp.ConnPool.GetOneAvailableConn()
	if err != nil {
		return nil, err
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to forward request: %v", err)
	}

	return resp, nil
}
func (msp *ModelServiceProxyImpl) StreamResponse(writer http.ResponseWriter, resp *http.Response) error {
	for k, v := range resp.Header {
		for _, vv := range v {
			writer.Header().Add(k, vv)
		}
	}
	writer.WriteHeader(resp.StatusCode)
	_, err := io.Copy(writer, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to stream response: %v", err)
	}
	return err
}
